import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
import os
from typing import List, Dict, Optional


class SimilarStockRecommender:
    def __init__(self, data_dir: str = "hs300_data"):
        """初始化推荐器，加载图模型并配置中文字体"""
        self.data_dir = data_dir
        self.G = None
        self.chinese_font = self._get_valid_chinese_font()  # 存储字体名称
        self.industries = []  # 所有行业列表
        self.industry_example_stocks = {}  # 各行业示例股票
        self.stock_name_to_codes = {}  # 股票名称→代码映射（用于名称查询）
        
        # 检查数据文件
        self._check_data_files()
        
        # 加载并构建异构图（同时初始化名称→代码映射）
        self._load_and_build_graph()
        
        # 获取各行业示例股票
        self._get_industry_example_stocks()
        
        # 验证图结构完整性
        self._validate_graph()
        
        # 显示行业示例股票
        self._display_industry_examples()

    def _get_valid_chinese_font(self) -> str:
        """检测并返回可用的中文字体名称"""
        candidate_fonts = [
            "SimHei",          # Windows 默认中文字体
            "Microsoft YaHei", # Windows 常用中文字体
            "PingFang SC",     # macOS 默认中文字体
            "Heiti SC",        # macOS 常用中文字体
            "WenQuanYi Zen Hei",# Linux 常用中文字体
        ]
        
        # 检测系统字体
        available_fonts = set([f.name for f in font_manager.fontManager.ttflist])
        for font in candidate_fonts:
            if font in available_fonts:
                print(f"✅ 检测到可用中文字体：{font}，将用于中文渲染")
                return font
        
        # 加载本地字体
        custom_font_path = os.path.join(self.data_dir, "fonts", "SimHei.ttf")
        if os.path.exists(custom_font_path):
            font_manager.fontManager.addfont(custom_font_path)
            font_manager._rebuild()
            if "SimHei" in [f.name for f in font_manager.fontManager.ttflist]:
                print(f"✅ 已成功加载本地字体：{custom_font_path}")
                return "SimHei"
            else:
                print(f"❌ 本地字体加载失败：{custom_font_path}（文件可能损坏）")
        
        # 无字体时报错
        raise RuntimeError(
            "❌ 无法找到可用中文字体！请按以下步骤解决：\n"
            "1. 下载 SimHei.ttf 字体文件\n"
            "2. 创建目录：hs300_data/fonts\n"
            "3. 将 SimHei.ttf 放入该目录后重试"
        )

    def _check_data_files(self) -> None:
        """检查必要的数据文件是否存在"""
        required_files = [
            "graph_stock_nodes.csv",
            "graph_stock_links.csv",
            "graph_industry_links.csv"
        ]
        for file in required_files:
            file_path = os.path.join(self.data_dir, file)
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"缺少数据文件：{file_path}")

    def _load_and_build_graph(self) -> None:
        """加载数据并构建异构图，同时初始化股票名称→代码映射"""
        # 加载股票节点数据
        stock_nodes = pd.read_csv(
            os.path.join(self.data_dir, "graph_stock_nodes.csv"),
            encoding="utf-8-sig"
        )
        
        # 加载股票-股票联动边
        stock_links = pd.read_csv(
            os.path.join(self.data_dir, "graph_stock_links.csv"),
            encoding="utf-8-sig"
        )
        
        # 加载股票-行业归属边
        industry_links = pd.read_csv(
            os.path.join(self.data_dir, "graph_industry_links.csv"),
            encoding="utf-8-sig"
        )
        
        # 保存行业列表
        self.industries = sorted(industry_links["industry"].unique().tolist())

        # 构建异构图
        self.G = nx.Graph()
        
        # 添加股票节点（含属性），并构建名称→代码映射
        for _, row in stock_nodes.iterrows():
            stock_code = row["stock_code"]
            stock_name = row["stock_name"].strip()  # 去除名称前后空格
            self.G.add_node(
                stock_code,
                type="stock",
                stock_name=stock_name,
                industry=row["industry"],
                volatility=row["volatility"],
                avg_turnover=row["avg_turnover_billion"]
            )
            # 构建名称→代码映射（支持模糊查询）
            if stock_name not in self.stock_name_to_codes:
                self.stock_name_to_codes[stock_name] = []
            self.stock_name_to_codes[stock_name].append(stock_code)
        
        # 添加行业节点
        industries = industry_links["industry"].unique()
        self.G.add_nodes_from(industries, type="industry")
        
        # 添加股票-股票联动边（带权重）
        for _, row in stock_links.iterrows():
            self.G.add_edge(
                row["stock1"], row["stock2"],
                type="stock_corr", weight=row["corr"]
            )
        
        # 添加股票-行业归属边
        for _, row in industry_links.iterrows():
            self.G.add_edge(
                row["stock_code"], row["industry"],
                type="stock_industry", weight=1.0
            )

    def _get_industry_example_stocks(self) -> None:
        """为每个行业选择一只代表性股票作为示例"""
        # 按行业分组，为每个行业选择评分最高的股票
        stock_nodes = pd.read_csv(
            os.path.join(self.data_dir, "graph_stock_nodes.csv"),
            encoding="utf-8-sig"
        )
        
        # 计算综合评分：成交额/波动率（值越高越好）
        stock_nodes["score"] = stock_nodes["avg_turnover_billion"] / (stock_nodes["volatility"] + 1e-8)
        
        # 按行业分组，取评分最高的股票作为示例
        for industry in self.industries:
            industry_stocks = stock_nodes[stock_nodes["industry"] == industry]
            if not industry_stocks.empty:
                top_stock = industry_stocks.sort_values(by="score", ascending=False).iloc[0]
                self.industry_example_stocks[industry] = {
                    "stock_code": top_stock["stock_code"],
                    "stock_name": top_stock["stock_name"],
                    "score": top_stock["score"]
                }

    def _display_industry_examples(self) -> None:
        """展示各行业的示例股票，引导用户输入"""
        print("\n📌 各行业参考股票示例（可输入代码或名称）：")
        # 控制每行显示的行业数量，避免输出过长
        stocks_per_line = 3
        for i, (industry, stock_info) in enumerate(self.industry_example_stocks.items()):
            line_end = "\n" if (i + 1) % stocks_per_line == 0 else "  |  "
            print(f"{industry}: {stock_info['stock_code']}（{stock_info['stock_name'][:6]}）", end=line_end)
        print("\n")

    def _validate_graph(self) -> None:
        """验证图结构并输出统计信息"""
        stock_cnt = len([n for n, d in self.G.nodes(data=True) if d["type"] == "stock"])
        link_cnt = len([e for e in self.G.edges(data=True) if e[2]["type"] == "stock_corr"])
        print(f"✅ 图模型加载完成：")
        print(f"   - 股票节点数：{stock_cnt}（沪深300成分股）")
        print(f"   - 股票-股票联动边数：{link_cnt}（带联动系数权重）")
        print(f"   - 支持输入：股票代码（如 sz.002230）或股票名称（如 科大讯飞）\n")

    def _find_similar_stocks(self, input_str: str) -> List[str]:
        """模糊匹配可能的股票代码或名称"""
        matches = []
        input_str_lower = input_str.lower()
        for node in self.G.nodes:
            if self.G.nodes[node]["type"] == "stock":
                # 匹配股票代码（包含输入字符串）
                if input_str_lower in node.lower():
                    matches.append(node)
                # 匹配股票名称（包含输入字符串）
                elif input_str_lower in self.G.nodes[node]["stock_name"].lower():
                    matches.append(node)
        return list(set(matches))  # 去重

    def _get_stock_code_by_name(self, input_name: str) -> Optional[str]:
        """
        根据股票名称模糊查询并获取股票代码
        返回：匹配到的股票代码（None表示无匹配或用户取消）
        """
        input_name_lower = input_name.lower().strip()
        # 1. 收集所有名称包含输入字符串的股票
        matched_stocks = []
        for stock_name, codes in self.stock_name_to_codes.items():
            if input_name_lower in stock_name.lower():
                for code in codes:
                    matched_stocks.append({
                        "code": code,
                        "name": stock_name,
                        "industry": self.G.nodes[code]["industry"]
                    })
        
        # 2. 处理匹配结果
        if not matched_stocks:
            print(f"❌ 未找到名称包含「{input_name}」的股票")
            return None
        
        # 去重（避免同一股票多次出现）
        unique_stocks = []
        seen_codes = set()
        for stock in matched_stocks:
            if stock["code"] not in seen_codes:
                unique_stocks.append(stock)
                seen_codes.add(stock["code"])
        
        # 3. 单个匹配结果：直接返回
        if len(unique_stocks) == 1:
            stock = unique_stocks[0]
            print(f"✅ 匹配到股票：{stock['code']}（{stock['name']} - {stock['industry']}）")
            return stock["code"]
        
        # 4. 多个匹配结果：让用户选择
        print(f"\n⚠️ 找到{len(unique_stocks)}只名称包含「{input_name}」的股票，请选择：")
        for idx, stock in enumerate(unique_stocks, 1):
            print(f"   {idx}. 代码：{stock['code']} | 名称：{stock['name']} | 行业：{stock['industry']}")
        print(f"   0. 取消选择")
        
        # 5. 处理用户选择
        while True:
            try:
                choice = input("\n请输入选择的序号（0-取消）：").strip()
                if choice == "0":
                    print("🔄 已取消选择")
                    return None
                choice_idx = int(choice) - 1
                if 0 <= choice_idx < len(unique_stocks):
                    selected_stock = unique_stocks[choice_idx]
                    print(f"✅ 您选择的股票：{selected_stock['code']}（{selected_stock['name']}）")
                    return selected_stock["code"]
                else:
                    print(f"❌ 无效序号，请输入1-{len(unique_stocks)}或0")
            except ValueError:
                print(f"❌ 输入无效，请输入数字序号")

    def recommend_similar(self, input_stock: str, top_k: int = 5) -> List[Dict]:
        """推荐与输入股票联动性强的股票（input_stock为股票代码）"""
        # 1. 验证输入是否为有效股票代码
        if input_stock not in self.G.nodes:
            similar_stocks = self._find_similar_stocks(input_stock)
            error_msg = f"股票[{input_stock}]不存在（需为沪深300成分股）"
            if similar_stocks:
                error_msg += f"\n可能的股票：{[f'{code}（{self.G.nodes[code]["stock_name"]}）' for code in similar_stocks[:5]]}"  # 最多显示5个
            return [{"error": error_msg}]
            
        if self.G.nodes[input_stock]["type"] != "stock":
            return [{"error": f"[{input_stock}]不是股票节点，请输入股票代码或名称"}]

        # 2. 提取联动股票候选
        similar_candidates = []
        for neighbor in self.G.neighbors(input_stock):
            edge_data = self.G[input_stock][neighbor]
            if (self.G.nodes[neighbor]["type"] == "stock") and (edge_data["type"] == "stock_corr"):
                similar_candidates.append({
                    "stock_code": neighbor,
                    "corr": edge_data["weight"],
                    "stock_name": self.G.nodes[neighbor]["stock_name"],
                    "industry": self.G.nodes[neighbor]["industry"],
                    "volatility": round(self.G.nodes[neighbor]["volatility"], 4),
                    "avg_turnover": round(self.G.nodes[neighbor]["avg_turnover"], 2)
                })

        if not similar_candidates:
            return [{"error": f"股票[{input_stock}]（{self.G.nodes[input_stock]['stock_name']}）无关联股票，无法推荐"}]

        # 3. 排序并生成结果
        similar_candidates.sort(key=lambda x: x["corr"], reverse=True)
        input_stock_name = self.G.nodes[input_stock]["stock_name"]
        result = []
        for idx, stock in enumerate(similar_candidates[:top_k], 1):
            result.append({
                "推荐排名": idx,
                "股票代码": stock["stock_code"],
                "股票名称": stock["stock_name"],
                "所属行业": stock["industry"],
                "与输入股票的联动系数": round(stock["corr"], 4),
                "波动率（风险）": stock["volatility"],
                "平均成交额（亿元）": stock["avg_turnover"],
                "推荐理由": f"基于异构图中“股票-股票”联动边：{stock['stock_name']}（{stock['stock_code']}）与{input_stock_name}（{input_stock}）的价格联动系数为{round(stock['corr'], 4)}，走势高度同步，适合追求同趋势布局的老手"
            })
        return result

    def evaluate_recommendation(self, test_stock_cnt: int = 5) -> None:
        """量化评估荐股结果的有效性"""
        print("\n=== 知股荐股量化评估 ===")
        test_stocks = [n for n, d in self.G.nodes(data=True) if d["type"] == "stock"][:test_stock_cnt]
        avg_corr_list = []
        
        for stock in test_stocks:
            recs = self.recommend_similar(input_stock=stock, top_k=5)
            if "error" not in recs[0]:
                avg_corr = sum(r["与输入股票的联动系数"] for r in recs) / len(recs)
                avg_corr_list.append(avg_corr)
                print(f"✅ 测试股票[{stock}]（{self.G.nodes[stock]['stock_name']}）：平均联动系数={avg_corr:.4f}")
        
        if avg_corr_list:
            overall_avg = sum(avg_corr_list) / len(avg_corr_list)
            print(f"\n📊 整体评估结果：{test_stock_cnt}只测试股票的平均联动系数={overall_avg:.4f}")
            print("✅ 评估结论：推荐结果联动性强（平均系数≥0.6），符合老手同趋势布局需求" if overall_avg >= 0.6 else "⚠️  建议优化数据质量")

    def visualize_recommendation(self, input_stock: str, top_k: int = 5) -> None:
        """可视化荐股路径，默认显示所有推荐的股票（top_k=5）"""
        recs = self.recommend_similar(input_stock=input_stock, top_k=top_k)
        if "error" in recs[0]:
            print(f"❌ 可视化失败：{recs[0]['error']}")
            return

        # 1. 全局字体配置
        plt.rcParams.update({
            "font.family": self.chinese_font,
            "axes.unicode_minus": False,
            "font.size": 8,  # 略微减小字体，避免标签重叠
            "figure.dpi": 100
        })

        # 2. 准备可视化节点 - 包含所有推荐股票
        rec_stock_codes = [r["股票代码"] for r in recs]
        visualize_nodes = [input_stock] + rec_stock_codes
        # 添加相关行业节点
        for node in visualize_nodes:
            if self.G.nodes[node]["type"] == "stock":
                visualize_nodes.append(self.G.nodes[node]["industry"])
        visualize_nodes = list(set(visualize_nodes))  # 去重
        subG = self.G.subgraph(visualize_nodes)

        # 3. 节点位置与样式 - 调整布局参数以适应更多节点
        fig, ax = plt.subplots(figsize=(14, 10), facecolor="white")  # 增大图表尺寸
        # 使用弹簧布局并调整参数，避免节点过于拥挤
        pos = nx.spring_layout(subG, seed=42, k=0.8, iterations=50)
        node_colors = []
        node_sizes = []
        for node in subG.nodes:
            if node == input_stock:
                node_colors.append("#FF6B6B")  # 输入股票：红色
                node_sizes.append(2000)
            elif self.G.nodes[node]["type"] == "stock":
                node_colors.append("#4ECDC4")  # 推荐股票：蓝色
                node_sizes.append(1500)
            else:
                node_colors.append("#96CEB4")  # 行业：绿色
                node_sizes.append(1200)

        # 4. 绘制边
        nx.draw(
            subG, pos, ax=ax,
            edgelist=subG.edges(),
            edge_color="#D3D3D3",
            width=1.2,
            with_labels=False
        )

        # 5. 绘制节点
        nx.draw(
            subG, pos, ax=ax,
            nodelist=subG.nodes(),
            node_color=node_colors,
            node_size=node_sizes,
            edgecolors="#333333",
            linewidths=1.0,
            with_labels=False
        )

        # 6. 绘制节点标签 - 缩短名称以避免重叠
        node_labels = {}
        for node in subG.nodes:
            if self.G.nodes[node]["type"] == "stock":
                stock_name = self.G.nodes[node]["stock_name"]
                # 进一步缩短名称显示（代码显示后6位，名称显示前4位）
                node_labels[node] = f"{node.split('.')[1]}\n{stock_name[:4]}..."
            else:
                node_labels[node] = node[:6] + "..." if len(node) > 6 else node

        nx.draw_networkx_labels(
            subG, pos, ax=ax,
            labels=node_labels,
            font_family=self.chinese_font,
            font_size=7,  # 减小标签字体
            font_weight="bold",
            font_color="#000000"
        )

        # 7. 绘制边权重标签
        edge_labels = {}
        for u, v, data in subG.edges(data=True):
            if data["type"] == "stock_corr":
                edge_labels[(u, v)] = f"corr:{round(data['weight'], 2)}"

        nx.draw_networkx_edge_labels(
            subG, pos, ax=ax,
            edge_labels=edge_labels,
            font_family=self.chinese_font,
            font_size=6,  # 减小边标签字体
            font_color="#FF6B6B",
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8)
        )

        # 8. 设置标题并保存
        input_stock_name = self.G.nodes[input_stock]["stock_name"]
        save_path = os.path.join(self.data_dir, f"similar_stock_{input_stock}_{input_stock_name}.png")
        
        ax.set_title(
            f"知股荐股路径可视化：{input_stock_name}（{input_stock}）→ 推荐股票",
            fontdict={"fontfamily": self.chinese_font, "fontsize": 12},
            pad=20
        )
        
        plt.subplots_adjust(left=0.05, right=0.95, top=0.9, bottom=0.1)
        plt.savefig(
            save_path,
            dpi=300,
            bbox_inches="tight",
            facecolor="white"
        )
        plt.close()
        print(f"\n📸 可视化结果已保存至：{save_path}")

    def run(self):
        """运行知股荐股主流程（支持代码/名称输入）"""
        print("="*60)
        print("📈 知股荐股系统（支持输入股票代码或名称）")
        print("="*60)
        
        while True:
            # 1. 让用户选择输入类型
            print("\n请选择输入方式：")
            print("   1. 输入股票代码（如：sz.002230）")
            print("   2. 输入股票名称（如：科大讯飞）")
            input_type_choice = input("请输入序号（1/2）：").strip()
            
            # 2. 处理输入类型选择
            input_stock_code = None
            if input_type_choice == "1":
                # 输入股票代码
                input_code = input("请输入股票代码：").strip()
                if not input_code:
                    print("❌ 输入不能为空，请重新输入")
                    continue
                # 验证代码并获取匹配结果
                similar_codes = self._find_similar_stocks(input_code)
                if not similar_codes:
                    print(f"❌ 未找到代码包含「{input_code}」的股票")
                    continue
                # 处理单个/多个代码匹配
                if len(similar_codes) == 1:
                    input_stock_code = similar_codes[0]
                else:
                    print(f"\n找到{len(similar_codes)}个匹配的股票代码，请选择：")
                    for idx, code in enumerate(similar_codes, 1):
                        print(f"   {idx}. {code}（{self.G.nodes[code]['stock_name']}）")
                    while True:
                        try:
                            choice = int(input("请输入选择的序号：").strip()) - 1
                            if 0 <= choice < len(similar_codes):
                                input_stock_code = similar_codes[choice]
                                break
                            else:
                                print(f"❌ 无效序号，请输入1-{len(similar_codes)}")
                        except ValueError:
                            print("❌ 输入无效，请输入数字序号")
            
            elif input_type_choice == "2":
                # 输入股票名称（调用名称查询方法）
                input_name = input("请输入股票名称：").strip()
                if not input_name:
                    print("❌ 输入不能为空，请重新输入")
                    continue
                input_stock_code = self._get_stock_code_by_name(input_name)
                if not input_stock_code:
                    continue  # 无匹配或用户取消，重新循环
            
            else:
                # 无效输入类型
                print("❌ 无效选择，请输入1或2")
                continue
            
            # 3. 执行推荐逻辑
            if input_stock_code:
                input_stock_name = self.G.nodes[input_stock_code]["stock_name"]
                print(f"\n=== 知股荐股结果（输入股票：{input_stock_code} - {input_stock_name}）===")
                recommendations = self.recommend_similar(input_stock=input_stock_code, top_k=5)
                
                # 显示推荐结果
                for rec in recommendations:
                    if "error" in rec:
                        print(f"❌ {rec['error']}")
                    else:
                        print(f"\n【第{rec['推荐排名']}名】")
                        print(f"股票代码：{rec['股票代码']} | 股票名称：{rec['股票名称']}")
                        print(f"所属行业：{rec['所属行业']} | 联动系数：{rec['与输入股票的联动系数']}")
                        print(f"波动率：{rec['波动率（风险）']} | 平均成交额：{rec['平均成交额（亿元）']}亿")
                        print(f"推荐理由：{rec['推荐理由']}")
                
                # 询问是否可视化
                visualize_choice = input("\n是否生成该股票与推荐股票的关系图？（输入y生成，其他键跳过）：").strip().lower()
                if visualize_choice == "y":
                    self.visualize_recommendation(input_stock=input_stock_code, top_k=5)
            
            # 4. 询问是否继续
            continue_choice = input("\n是否继续查询其他股票？（输入y继续，其他键退出）：").strip().lower()
            if continue_choice != "y":
                print("\n👋 感谢使用知股荐股系统，程序退出！")
                print("="*60)
                break


# 程序入口
if __name__ == "__main__":
    try:
        recommender = SimilarStockRecommender(data_dir="hs300_data")
        recommender.run()
    except Exception as e:
        print(f"\n程序运行出错：{e}")